#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project:     rest_flask
# @File:        ner_matcher.py
# @Author:      Fan GuiChuan
# @Date:        2025/6/30 14:36
# @Version:     Python3.7
# @Description:
import os.path

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

CUR_FILE_PATH = os.path.realpath(__file__)
CUR_FILE_DIR = os.path.dirname(CUR_FILE_PATH)


class NerMatcher:
    def __init__(self, schema, ner_tag_id_map, model='ner_struct_bert_md'):
        model_path = os.path.join(os.environ['root_path'], f"models/{model}")
        self.ner_pip = pipeline(Tasks.siamese_uie, model_path, device='cpu')
        self.schema = schema
        self.ner_tag_id_map = ner_tag_id_map

    def set_schema(self, schema, ner_tag_id_map):
        self.schema = schema
        self.ner_tag_id_map = ner_tag_id_map

    def match_all(self, data):
        if self.ner_pip is None:
            return []
        pip_result = self.ner_pip(data, schema=self.schema)

        output = pip_result['output']
        output = output[0] if output else []
        result = []
        for tag in output:
            rule_info = self.ner_tag_id_map.get(tag['type'], {})
            result.append({
                'feature_id': rule_info.get('label_id'),
                'feature_name': rule_info.get('label_name'),
                'level_code': rule_info.get('level_code'),
                'level_id': rule_info.get('level_id'),
                'rule_id': rule_info.get('rule_id'),
                'template_id': rule_info.get('template_id'),
                'from': tag['offset'][0],
                'to': tag['offset'][1],
                'match': tag['span'],  # 提取匹配文本
                'identify_type_name': 'ner',
                'identify_type': 1,
            })

        return result


if __name__ == '__main__':
    ner_matcher = NerMatcher(
        {
            '疾病名称': None,
            '问好': None
        },
        {
            '疾病名称': {},
            '问好': {
                'label_name': '问好'
            }
         },
        '问好（6）'
    )
    result = ner_matcher.match_all('你们好')

    print(result)